跳到主要内容

PyTorch 中的延后初始化

到目前为止,我们忽略了建立网络时需要做的以下这些事情:

  • 我们定义了网络架构,但没有指定输入维度。
  • 我们添加层时没有指定前一层的输出维度。
  • 我们在初始化参数时,甚至没有足够的信息来确定模型应该包含多少参数。

有些读者可能会对我们的代码能运行感到惊讶。 毕竟,深度学习框架无法判断网络的输入维度是什么。 这里的诀窍是框架的延后初始化(defers initialization), 即直到数据第一次通过模型传递时,框架才会动态地推断出每个层的大小。

在深度学习框架中,延后初始化(也称为懒初始化)是一种策略,其中模型的参数在声明时不立即初始化,而是等待输入数据的形状确定后再进行初始化。这种策略的主要优点是它允许用户在不明确指定输入形状的情况下定义模型,从而提供更大的灵活性。

PyTorch的行为与此略有不同。在PyTorch中,当你定义一个模块或层(例如nn.Linearnn.Conv2d)时,你通常需要指定输入和输出的大小或特征数。这意味着在大多数情况下,PyTorch会在模块或层被定义时立即初始化其参数。

然而,对于某些特定的模块或自定义层,你可能希望实现自己的延后初始化逻辑。这可以通过以下步骤实现:

  1. __init__ 方法中,不立即调用参数的初始化函数。
  2. forward 方法中,检查参数是否已经被初始化。如果没有,则根据输入数据的形状进行初始化。

以下是一个简单的示例,展示如何在自定义的全连接层中实现延后初始化:

import torch.nn as nn
import torch.nn.functional as F

class LazyLinear(nn.Module):
def __init__(self, out_features, bias=True):
super(LazyLinear, self).__init__()
self.out_features = out_features
self.bias = bias
self.weight = None
if bias:
self.bias_tensor = None

def forward(self, x):
# 如果权重未被初始化
if self.weight is None:
in_features = x.size(-1)
self.weight = nn.Parameter(torch.randn(in_features, self.out_features))
if self.bias:
self.bias_tensor = nn.Parameter(torch.randn(self.out_features))

return F.linear(x, self.weight, self.bias_tensor)

# 使用
x = torch.randn(5, 3)
model = LazyLinear(10)
output = model(x)

在上述示例中,LazyLinear层在定义时不需要输入特征数。只有在第一次调用forward方法时,根据输入x的形状,权重才会被初始化。

总的来说,虽然PyTorch的大多数内置模块和层在定义时就进行了初始化,但你可以在自定义模块或层中实现延后初始化的逻辑,以提供更大的灵活性。